package com.chb.dmp.utils

import com.chb.dmp.conf.Constant
import org.apache.kudu.Schema
import org.apache.kudu.client.CreateTableOptions
import org.apache.kudu.spark.kudu.KuduContext
import org.apache.spark.sql.{DataFrame, Dataset, SaveMode, SparkSession}

class KuduHelper {
  private var kuduContext: KuduContext = _
  private var spark: SparkSession = _
  private var ds: Dataset[_] = _

  def this(spark: SparkSession) {
    this()
    this.spark = spark
    val KUDU_MASTER = Constant.KUDU_MASTER
    kuduContext = new KuduContext(KUDU_MASTER, spark.sparkContext)
  }

  def this(ds: Dataset[_]) {
    this()
    this.ds = ds
  }

  // 二： 具体功能的开发
  /*
   * 2.1、创建表
   *    1、通过隐式转换，将SparkSession 转为 KuduHelper
   *    2、通过KuduHelper中的createKuduTable 就可可以创建一张表
   * 调用方式： spark.createKuduTable()
   */
  def createKuduTable(tableName: String, schema: Schema, options: CreateTableOptions): Unit = {
    // 通过kuduContext 创建表
    if (kuduContext.tableExists(tableName)) kuduContext.deleteTable(tableName)

    //
    import scala.collection.JavaConverters._
    //    val optionsI = new CreateTableOptions()
    //      .setNumReplicas(1)
    //      .addHashPartitions(Seq("KEY").asJava, 2)

    kuduContext.createTable(tableName, schema, options)
  }

  /**
   *  2.2、读取表
   */
  def readKuduTable(tableName: String): Option[DataFrame] = {
    val KUDU_MASTER = Constant.KUDU_MASTER

    //定义 map 集合，封装 kudu 的 master 地址和要读取的表名
    val options = Map(
      "kudu.master" -> KUDU_MASTER,
      "kudu.table" -> tableName
    )
    if (kuduContext.tableExists(tableName)) {
      val result = spark.read
        .options(options)
        .format("kudu").load()
      Some(result)
    } else {
      None
    }
  }

  /**
   * 2.3、写入数据
   *
   * 调用 dataset.saveToKudu
   */
  def saveToKudu(tableName: String): Unit = {
    if (ds == null) {
      throw new RuntimeException("请从 DataSet上开始调用")
    }
    val KUDU_MASTER = Constant.KUDU_MASTER
    val options = Map(
      "kudu.master" -> KUDU_MASTER,
      "kudu.table" -> tableName
    )

    import org.apache.kudu.spark.kudu._
    ds.write
      .options(options)
      .mode(SaveMode.Append)
      .kudu
  }

}

object KuduHelper {
  // 一、隐式转换， SparkSession -> KuduHelper DataFrame -> KuduHelper
  // 就是将SparkHelper中的方法和功能 扩展到 SparkSession中

  /**
   * 设计隐式转换的时候， 只需要考虑一件事，把XX转为YY
   * XX 是SparkSession, 转换函数的传入参数是SparkSession
   * YY 是KuduHelper, 转换函数的结果类型就是KuduHelper
   */
  implicit def SparkSessionToKuduHelper(spark: SparkSession): KuduHelper = {
    new KuduHelper(spark)
  }

  /**
   * DataFrame -> KuduHelper
   *
   * @param ds // 此处不可以使用Dataset[Any], 如果想所有的Dataset都可以使用这个隐式转换， 必须使用Dataset[_],使用通配符
   * @return
   */
  implicit def DFToKudu(ds: Dataset[_]): KuduHelper = {
    new KuduHelper(ds)
  }

}
